# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import aiofiles
import importlib
from aiofiles import os as aio_os
from typing import List, Dict
from datetime import datetime
from fastapi import FastAPI
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.date import DateTrigger
from common.cache.queue_cache import QueueCache
from common.enums.apply import MusicEnum
from common.enums.apply import VideoEnum
from common.enums.apply import DrawsEnum
from common.enums.agent import AgentEnum
from common.models.sys import SysCrontabModel
from common.models.music import MusicRecordsModel
from common.models.video import VideoRecordsModel
from common.models.draws import DrawsRecordsModel
from common.models.agents import AgentKnowSplittingModel
from common.postgres.public import PgKnowledgeModel

scheduler = AsyncIOScheduler()


class AppEvents:
    @classmethod
    async def startup(cls, _app: FastAPI):
        scheduler.add_job(cls._inject_crontab, DateTrigger(run_date=datetime.now()))
        scheduler.start()

    @classmethod
    async def shutdown(cls, _app: FastAPI):
        scheduler.shutdown()

    @classmethod
    async def _inject_crontab(cls):
        crontab_tips = ""
        crontab_lists = await SysCrontabModel.filter(is_delete=0).order_by("-id").all()
        for crontab in crontab_lists:
            s = {1: "Success", 2: "Stop", 3: "Error"}
            if crontab.status != 1:
                e = s.get(crontab.status)
                crontab_tips += f"| {crontab.name} | {crontab.command} | {crontab.concurrent}    | {e} \n"
                continue

            try:
                module = importlib.import_module(crontab.command)
            except ModuleNotFoundError:
                raise Exception(f"The scheduled task module does not exist: {crontab.command}")

            func = getattr(module, "execute", None)
            if not func:
                raise Exception(f"Task execution method does not exist: {crontab.command}")

            # 获取任务参数
            params: Dict[str, any] = json.loads(crontab.params or "{}") if crontab.params else {}

            # 获取任务规则
            rules: List[dict] = json.loads(crontab.rules)

            # 处理触发条件
            condition = {}
            for item in rules:
                if crontab.trigger == "interval" and item.get("key") not in ["start_date", "end_date"]:
                    condition[item.get("key")] = int(item.get("value"))
                else:
                    condition[item.get("key")] = item.get("value")

            # 配置触发条件
            _trigger_fun: any = None
            if crontab.trigger == "interval":
                _trigger_fun = IntervalTrigger(**condition)
            elif crontab.trigger == "cron":
                _trigger_fun = CronTrigger(**condition)
            elif crontab.trigger == "date":
                _trigger_fun = DateTrigger(**condition)

            # 加入到任务中
            for i in range(crontab.concurrent):
                job = crontab.command + "." + str(i + 1)
                params["w_id"] = crontab.id
                params["w_ix"] = int(i + 1)
                params["w_job"] = job
                scheduler.add_job(func, _trigger_fun, id=job, name=crontab.name, kwargs=params)

            # 启动成功提示
            e = s.get(crontab.status)
            crontab_tips += f"| {crontab.name} | {crontab.command} | {crontab.concurrent}    | {e}  \n"

        await cls._reset_queue()

        if await aio_os.path.exists("./banner.txt"):
            async with aiofiles.open("./banner.txt", mode="r", encoding="utf-8") as f:
                content = await f.read()
            crontab_tips = crontab_tips if crontab_tips else "|"
            print(content.replace("{{crontab}}", crontab_tips.rstrip("\n")))

    @classmethod
    async def _reset_queue(cls):
        # 向量队列重置
        await QueueCache.delete_queue(QueueCache.EM_JOB)
        embed_lists = await (PgKnowledgeModel
                             .filter(is_delete=0)
                             .filter(status__in=[AgentEnum.EMB_LINE, AgentEnum.EMB_ING])
                             .order_by("create_time", "-status")
                             .values("uuid"))
        for item in embed_lists:
            await QueueCache.push_em(str(item["uuid"]))

        # 拆分队列重置
        await QueueCache.delete_queue(QueueCache.QA_JOB)
        split_lists = await (AgentKnowSplittingModel
                             .filter(status__in=[AgentEnum.EMB_LINE, AgentEnum.EMB_ING])
                             .order_by("id", "-status")
                             .values("id"))
        for item in split_lists:
            await QueueCache.push_qa(str(item["id"]))

        # 绘画队列重置
        await QueueCache.delete_queue(QueueCache.DRAW_JOB)
        draws_lists = await (DrawsRecordsModel.filter(is_delete=0, status=DrawsEnum.STATUS_ING).values("id"))
        for item in draws_lists:
            await QueueCache.push_draw(item["id"])

        # 视频队列重置
        await QueueCache.delete_queue(QueueCache.VIDEO_JOB)
        video_lists = await (VideoRecordsModel.filter(is_delete=0, status=VideoEnum.STATUS_ING).values("id"))
        for item in video_lists:
            await QueueCache.push_video(item["id"])

        # 音乐队列重置
        await QueueCache.delete_queue(QueueCache.VIDEO_JOB)
        music_lists = await (MusicRecordsModel.filter(is_delete=0, status=MusicEnum.STATUS_ING).values("id"))
        for item in music_lists:
            await QueueCache.push_music(item["id"])
